package com.fuyun.rpc.common.security;

import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * 对称密钥认证实现
 */
public class SymmetricAuthenticator implements Authenticator {

	private static final String HMAC_ALGORITHM = "HmacSHA256";
	private final Map<String, String> clientSecrets = new HashMap<>();

	public SymmetricAuthenticator() {
		// 初始化一些测试客户端凭证
		clientSecrets.put("test-client", "test-secret");
	}

	@Override
	public boolean authenticate(ClientCredentials credentials) {
		String secret = clientSecrets.get(credentials.getClientId());
		if (secret == null) {
			return false;
		}

		try {
			String expectedSignature = generateSignature(credentials, secret);
			return expectedSignature.equals(credentials.getSignature());
		} catch (Exception e) {
			return false;
		}
	}

	@Override
	public String generateToken(ServiceInfo serviceInfo) {
		try {
			String data =
				serviceInfo.getServiceName() + ":" + serviceInfo.getServiceVersion() + ":" + serviceInfo.getIp() + ":" + serviceInfo.getPort();
			return generateHMAC(data, "rpc-token-secret");
		} catch (Exception e) {
			throw new RuntimeException("生成令牌失败", e);
		}
	}

	@Override
	public boolean validateToken(String token, ServiceInfo serviceInfo) {
		try {
			String expectedToken = generateToken(serviceInfo);
			return expectedToken.equals(token);
		} catch (Exception e) {
			return false;
		}
	}

	protected String generateSignature(ClientCredentials credentials, String secret)
		throws NoSuchAlgorithmException, InvalidKeyException {
		String data = credentials.getClientId() + ":" + credentials.getTimestamp();
		return generateHMAC(data, secret);
	}

	private String generateHMAC(String data, String secret) throws NoSuchAlgorithmException, InvalidKeyException {
		Mac mac = Mac.getInstance(HMAC_ALGORITHM);
		SecretKeySpec secretKeySpec = new SecretKeySpec(secret.getBytes(), HMAC_ALGORITHM);
		mac.init(secretKeySpec);
		byte[] rawHmac = mac.doFinal(data.getBytes());
		return Base64.getEncoder()
			.encodeToString(rawHmac);
	}
}